#include <chrono>
#include <fstream>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <string>
#include <thread>
#include <vector>

#include "common.hpp"
#include "det/postprocess.h"
#include "det/preprocess.h"
#include "det/rkpt.hpp"

// 测试函数：验证模型加载和基本推理
void test_model_loading(const std::string &model_path) {
    std::cout << "=== 测试模型加载 ===" << std::endl;
    
    // 检查模型文件是否存在
    std::ifstream model_file(model_path);
    if (!model_file.good()) {
        std::cerr << "错误：模型文件不存在: " << model_path << std::endl;
        return;
    }
    model_file.close();
    
    std::cout << "模型文件存在: " << model_path << std::endl;
    
    // 尝试初始化模型
    RkPt *rknn_model = new RkPt(model_path);
    if (rknn_model->init(nullptr, false) != 0) {
        std::cerr << "RKNN模型初始化失败" << std::endl;
        delete rknn_model;
        return;
    }
    
    std::cout << "✓ RKNN模型初始化成功" << std::endl;
    
    // 创建测试图像
    cv::Mat test_img(480, 640, CV_8UC3, cv::Scalar(128, 128, 128));
    
    // 在图像中心画一个矩形作为测试目标
    cv::rectangle(test_img, cv::Point(200, 150), cv::Point(440, 330), cv::Scalar(0, 255, 0), 2);
    cv::putText(test_img, "Test Object", cv::Point(220, 140), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 255, 0), 2);
    
    std::cout << "=== 测试推理 ===" << std::endl;
    
    // 执行推理
    DetectionResultsGroup results = rknn_model->infer(test_img, 0);
    
    std::cout << "推理完成，检测到 " << results.dets.size() << " 个目标" << std::endl;
    
    // 显示结果
    for (size_t i = 0; i < results.dets.size(); i++) {
        const auto &det = results.dets[i];
        std::cout << "目标 " << i + 1 << ": " << det.det_name 
                  << " 置信度: " << det.score 
                  << " 位置: " << det.box << std::endl;
    }
    
    // 保存测试图像
    cv::imwrite("test_input.jpg", test_img);
    std::cout << "测试图像已保存为 test_input.jpg" << std::endl;
    
    // 在结果图像上绘制检测框
    cv::Mat result_img = test_img.clone();
    for (size_t i = 0; i < results.dets.size(); i++) {
        const auto &det = results.dets[i];
        cv::rectangle(result_img, det.box, cv::Scalar(0, 0, 255), 2);
        std::string label = det.det_name + " " + std::to_string(det.score).substr(0, 4);
        cv::putText(result_img, label, cv::Point(det.box.x, det.box.y - 10), 
                    cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255), 2);
    }
    
    cv::imwrite("test_result.jpg", result_img);
    std::cout << "结果图像已保存为 test_result.jpg" << std::endl;
    
    delete rknn_model;
    std::cout << "=== 测试完成 ===" << std::endl;
}

int main(int argc, char *argv[]) {
    std::string model_path = "../model/wuzi.rknn"; // 默认模型路径
    
    // 解析命令行参数
    if (argc > 1) {
        model_path = argv[1];
    }
    
    std::cout << "使用模型: " << model_path << std::endl;
    
    try {
        test_model_loading(model_path);
    }
    catch (const std::exception &e) {
        std::cerr << "错误: " << e.what() << std::endl;
        return -1;
    }
    
    return 0;
} 